/**
 \file CollinsKt.h
 Declaration of class CollinsKt.
 
 \author Thomas Burton 
 \date 2/28/12
 */

#ifndef _CollinsKt_H_
#define _CollinsKt_H_

#include <list>
#include <map>
#include <memory>
#include <set>
#include <sstream>
#include <string>
#include <utility>
#include <vector>

#ifndef __CINT__
#include <boost/shared_ptr.hpp>
#endif

#include <TAxis.h>

#include "tppmc/Bins.h"
#include "tppmc/DssFragmentation.h"
#include "tppmc/Enumerations.h"
#include "tppmc/FragmentationFunction.h"
#include "tppmc/Particle.h"
#include "tppmc/TF2J.h"

namespace tppmc {

namespace collins {

/**
 \todo Move implementation to .cxx file
 */
class FragBin {
   
public:
   
   int favoured; ///< favoured, unfavoured, or the average of those
   double depolMin;
   double depolMax;
   
   bool operator<(const FragBin& other) const {
      if(favoured < other.favoured) return true;
      if(favoured > other.favoured) return false;
      // favoured is equal, compare on depol range
      if(depolMin < other.depolMin) return true;
      if(depolMax > other.depolMax) return false;
      return depolMax < other.depolMax;
   }
   void Speak() const {
      std::cout << favoured << " " << depolMin << " " << depolMax << std::endl;
   }
   bool Match(const Hadron& hadron) const {
      int fav = 0; // By default assume unfavoured
      int htype = hadron.type;
      int ptype = hadron.mParton.type;
      switch(htype) {
         case PIPLUS:
            if(2 == ptype or -1 == ptype) fav = 1;
            break;
         case PIMINUS:
            if(1 == ptype or -2 == ptype) fav = 1;
            break;
         case PIZERO: // pi0 from u(bar), d(bar) is average --> fav = 2
            if(2 == ptype or -1 == ptype or
               -2 == ptype or 1 == ptype) fav = 2;
            break;
      } // switch
      if(fav not_eq favoured) return false;
      if(hadron.mParton.cos_spin_depol < depolMin or
         hadron.mParton.cos_spin_depol >= depolMax) return false;
      return true;
   }               
};

class FunctorNew; // Forward declaration so FunctionTable can create
                  // a pointer list
/**
 The collection of ROOT functions used by PositivityKt.
 The function definitions are all the same, but have different parameters.
 A different copy of the function is maintained for each permutation of
 parameters to avoid recomputation of the function integral each time the
 parameters change, which gets computationally expensive.
 */
class FunctionTable {
public:
   
   FunctionTable();
   
#ifndef __CINT__
   typedef boost::shared_ptr<tppmc::FragmentationFunction> FragFcnPtr;
   FunctionTable(FragFcnPtr);
#endif
   
   /** Destructor */
   virtual ~FunctionTable();
   
   /**
    Returns the ROOT wrapper around the fragmentation function
    for computing fragmentation for a hadron at this scale.
    The hadron's mParton attribute must be set and define type and
    cos_spin_depol.
    Returns NULL if there is no function for this parton/hadron and Q bin.
    The pointer is non-const because methods like GetRandom() and
    Integral() are non-const, but don't modify the function e.g. by
    changing parameters.
    */
   TF2* Find(const Hadron&, double /*Q2*/);
   
   /**
    Populate the table with functions for each bin in depol and Q.
    Functions are evaluated at the lower edges of the bins.
    */
   int Initialise(const Bins* depol, const Bins* Q);
   
   void Speak();
   
protected:
#ifndef __CINT__
   TAxis* mQ;
   typedef std::vector<TF2*> TF2Vec;
   typedef std::map<FragBin, TF2Vec> Table;
   std::map<FragBin, TF2Vec> mFunctions;
   TF2* BuildFunction(FragBin, double Q);
   FragFcnPtr mUnpol;
   std::list<FunctorNew*> mFunctors;
#endif
   ClassDef(tppmc::collins::FunctionTable, 1)
};


/**
 kT-dependent positivity bound based on a kT-integrated unpolarised
 fragmentation function.
 Only pion production is currently supported.
 The positivity bound is defined via
 |collins(z, pt)| = 2 * unpolarised(z, pt)
 where collins(z, pt) is the Anselmino et al. definition of the
 Collins function, such that
 fftotal(z, pt) = unpolarised(z, pt) + 0.5s.(q x pt)collins(z, pt)
 where s is the parton spin and q the parton 3-momentum.
 As the pt dependence is not well known, it is factorised according to
 unpolarised(z, pt) = unpolarised(z) * exp(-pt2/<pt2>) / pi / <pt2>
 pt2 is the squared hadron pt wrt parton 3-momentum direction and
 <pt2> is the mean pt2. This is settable as a parameter (default is
 0.2 GeV^2). The unpolarised fragmentation can be set via the constructor.
 The default is leading order DSS.
 \todo Move headers for code in collins namespace to tppmc/collins/
 subdirectory, but check impact on LinkDef/CINT dictionary first.
 */
class PositivityKt : public FragmentationFunction {
   
public:
   
   /**
    Initialises with default distribution (LO DSS) and mean pt2 = 0.2
    */
   PositivityKt();
   
#ifndef __CINT__
   PositivityKt(boost::shared_ptr<FragmentationFunction>,
                double meanPt2,
                const std::set<int>& hadrons = std::set<int>());
#endif
   
   /**
    Destructor.
    */
   virtual ~PositivityKt();
   
   /**
    Perform initialisation operations.
    
    Must be called before anything else.
    */
   virtual bool Initialise();
   
   /**
    Fragment the input parton to a single hadron at the requested scale.
    If the parton species is not supported, fragment via the normal
    DSS fragmentation function, with a random kT according to a
    Gaussian distribution.
    */
   virtual std::vector<tppmc::Hadron> Fragment(const Parton&,
                                               double scale);
   
   /**
    Evaluate the fragmentation function for the input parton to
    fragment to the input hadron at the requested scale.
    */
   virtual double Evaluate(const Parton&, const Hadron&, double scale);
   
   /**
    Returns true if the fragmentation of the input parton type
    is supported by this function.
    */
   virtual bool Supports(const Parton&) const;
   
   /**
    Returns true if the fragmentation of the input parton type
    is supported by this function.
    */
   virtual bool Supports(const Hadron&) const;
   
   /**
    Returns the name of this fragmentation function
    */
   virtual std::string Name() const;
   
   virtual void SetHadron(int = INVALID_HADRON);
   
   virtual void SetZRange(double lower, double upper);
   
   /**
    Lists hadron species supported by the fragmentation function.
    */
   virtual const std::vector<int>& Hadrons() const;
   
   /**
    Lists parton species supported by the fragmentation function.
    */
   virtual const std::vector<int>& Partons() const;
   
   virtual int ThrowHadronType(const Parton& parton, double Q2);
   
   /** Returns the minimum z for which the distribution is valid */
   virtual double ZMin() const;
   
   /** Returns the maximum z for which the distribution is valid */
   virtual double ZMax() const;
   
   /**
    The lowest fragmentation scale allowed by the fragmentation function.
    */
   virtual double MinimumScale() const;
   
   /**
    The highest fragmentation scale allowed by the fragmentation function.
    */
   virtual double MaximumScale() const;
   
   virtual void Qa(std::string filename);
   
protected:
   
   void ComputeFragmentationProbabilities();
   double mMeanPt2;
   
   std::vector<int> mPartons; ///< Supported parton PDG codes
   std::vector<int> mHadrons; ///< Supported hadron PDG codes
#ifndef __CINT__
   typedef boost::shared_ptr<FragmentationFunction> FragFcnPtr;
   FragFcnPtr mUnpol; //!< Unpolarised
                      ///< frag function
   
   typedef boost::shared_ptr<FragmentationTable> FragTablePtr;
   std::map<int, FragTablePtr> mProbabilities; //!< Probability for each
                                               ///< parton to fragment
                                               ///< each hadron at each Q
   /**\todo change to smart pointer*/
   std::auto_ptr<Bins> mQBins; //!< Bins in Q: NOT Q-squared!
   
   FunctionTable mFunctions;//!
#endif
   ClassDef(tppmc::collins::PositivityKt, 1)
};

/**
 \todo Move to its own file
 */
class KtDependent {
   
public:
   
   /**
    Information describing parameters 
    */
   class FragInfo {
      
   public:
      
      int favoured;
      double depolMin;
      double depolMax;
      
      bool operator<(const FragInfo& other) const {
         if(favoured < other.favoured) return true;
         if(favoured > other.favoured) return false;
         // favoured is equal, compare on depol range
         if(depolMin < other.depolMin) return true;
         if(depolMax > other.depolMax) return false;
         return depolMax < other.depolMax;
      }
      void Speak() const {
         std::cout << favoured << " " << depolMin << " " << depolMax << std::endl;
      }
      bool Match(const Hadron& hadron) const {
         int fav = -1;
         int htype = hadron.type;
         int ptype = hadron.mParton.type;
         switch(htype) {
            case PIPLUS:
               if(2 == ptype or -1 == ptype) fav = 1;
               else fav = 0;
               break;
            case PIMINUS:
               if(1 == ptype or -2 == ptype) fav = 1;
               else fav = 0;
               break;
            case PIZERO:
               if(2 == ptype or -1 == ptype or
                  -2 == ptype or 1 == ptype) fav = 1;
               else fav = 0;
               break;
         } // switch
         std::cout << "Matching: parton " << ptype << " hadron " << htype << " fav = " << fav << std::endl;
         if(fav not_eq favoured) return false;
         
         if(hadron.mParton.cos_spin_depol < depolMin or
            hadron.mParton.cos_spin_depol >= depolMax) return false;
         return true;
      }               
      
      ClassDef(tppmc::collins::KtDependent::FragInfo, 1)
   };
   
   
   typedef std::map<FragInfo, std::string> Table;
   Table mTable;
   
   KtDependent() {
      
      double depol[11] = {
         0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.
      };
      std::stringstream ss;
      for(int i(0); i < 2; ++i) {
         for(int j(0); j < 10; ++j) {
            FragInfo f;
            f.favoured = i;
            f.depolMin = depol[j];
            f.depolMax = depol[j + 1];
            ss.str(""); ss.clear();
            ss << "Favoured " << f.favoured << " dmin = " <<
            f.depolMin << " dmax = " << f.depolMax;
            mTable.insert(std::make_pair(f, ss.str()));
         } // for
      } // for
   }
   
   std::string Find(const Hadron& hadron) const;
   
   ClassDef(tppmc::collins::KtDependent, 1)
};

} // namespace collins

} // namespace tppmc

#endif
